#!/bin/bash
# -*- mode: sh; indent-tabs-mode: nil; sh-basic-offset: 2; -*-
# vim: et sts=2 sw=2
#
# 用于创建、修复或修改 SteamOS 安装的函数集合。
#

set -eu

die() { echo >&2 "!! $*"; exit 1; }
readvar() { IFS= read -r -d '' "$1" || true; }

DISK=/dev/nvme0n1
DISK_SUFFIX=p
DOPARTVERIFY=1

# 如果此目录存在，使用此目录中的 jupiter-biosupdate 二进制文件，
# 并在调用时将 JUPITER_BIOS_DIR 设置为此目录。
# 用于包含比基础镜像更新的 BIOS 负载。
VENDORED_BIOS_UPDATE=/home/deck/jupiter-bios
# 如果此目录存在，使用此目录中的 jupiter-controller-update 二进制文件，
# 并在调用时将 JUPITER_CONTROLLER_UPDATE_FIRMWARE_DIR 设置为此目录。
# 用于包含比基础镜像更新的控制器负载。
VENDORED_CONTROLLER_UPDATE=/home/deck/jupiter-controller-fw

# 分区表，sfdisk 格式，%%DISKPART%% 将被填入
PART_SIZE_ESP="256"
PART_SIZE_EFI="64"
PART_SIZE_ROOT="5120" # 这应该与输入磁盘的大小匹配
PART_SIZE_VAR="256"
PART_SIZE_HOME="100"

# 总大小 + 开始/结束处的 1MiB 填充用于 GPT 结构。
DISK_SIZE=$(( 2 + PART_SIZE_HOME + PART_SIZE_ESP + 2 * ( PART_SIZE_EFI + PART_SIZE_ROOT + PART_SIZE_VAR ) ))
# 对齐：使用 MiB 等通用大小且没有明确的起始偏移点会导致 sfdisk 默认对齐到 MiB 边界

TARGET_SECTOR_SIZE=512 # 传递给 `losetup` 进行模拟，影响 sfdisk 最终写入的扇区偏移。
readvar PARTITION_TABLE <<END_PARTITION_TABLE
  label: gpt
  %%DISKPART%%1: name="esp",      size=         ${PART_SIZE_ESP}MiB,  type=C12A7328-F81F-11D2-BA4B-00A0C93EC93B
  %%DISKPART%%2: name="efi-A",    size=         ${PART_SIZE_EFI}MiB,  type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7
  %%DISKPART%%3: name="efi-B",    size=         ${PART_SIZE_EFI}MiB,  type=EBD0A0A2-B9E5-4433-87C0-68B6B72699C7
  %%DISKPART%%4: name="rootfs-A", size=         ${PART_SIZE_ROOT}MiB, type=4F68BCE3-E8CD-4DB1-96E7-FBCAF984B709
  %%DISKPART%%5: name="rootfs-B", size=         ${PART_SIZE_ROOT}MiB, type=4F68BCE3-E8CD-4DB1-96E7-FBCAF984B709
  %%DISKPART%%6: name="var-A",    size=         ${PART_SIZE_VAR}MiB,  type=4D21B016-B534-45C2-A9FB-5C16E091FD2D
  %%DISKPART%%7: name="var-B",    size=         ${PART_SIZE_VAR}MiB,  type=4D21B016-B534-45C2-A9FB-5C16E091FD2D
  %%DISKPART%%8: name="home",     size=         ${PART_SIZE_HOME}MiB, type=933AC7E1-2EB4-4F13-B844-0E14E2AEF915
END_PARTITION_TABLE

# 理想目标设备上的分区号
FS_ESP=1
FS_EFI_A=2
FS_EFI_B=3
FS_ROOT_A=4
FS_ROOT_B=5
FS_VAR_A=6
FS_VAR_B=7
FS_HOME=8

diskpart() { echo "$DISK$DISK_SUFFIX$1"; }

##
## 工具颜色等
##

err() {
  echo >&2
  eerr "发生成像错误，请查看上方信息并重新启动进程。"
  sleep infinity
}
trap err ERR

_sh_c_colors=0
[[ -n $TERM && -t 1 && ${TERM,,} != dumb ]] && _sh_c_colors="$(tput colors 2>/dev/null || echo 0)"
sh_c() { [[ $_sh_c_colors -le 0 ]] || ( IFS=\; && echo -n $'\e['"${*:-0}m"; ); }

sh_quote() { echo "${@@Q}"; }
estat()    { echo >&2 "$(sh_c 32 1)::$(sh_c) $*"; }
emsg()     { echo >&2 "$(sh_c 34 1)::$(sh_c) $*"; }
ewarn()    { echo >&2 "$(sh_c 33 1);;$(sh_c) $*"; }
einfo()    { echo >&2 "$(sh_c 30 1)::$(sh_c) $*"; }
eerr()     { echo >&2 "$(sh_c 31 1)!!$(sh_c) $*"; }
die() { local msg="$*"; [[ -n $msg ]] || msg="脚本终止"; eerr "$msg"; exit 1; }
showcmd() { showcmd_unquoted "${@@Q}"; }
showcmd_unquoted() { echo >&2 "$(sh_c 30 1)+$(sh_c) $*"; }
cmd() { showcmd "$@"; "$@"; }

# 格式化辅助函数
fmt_ext4()  { [[ $# -eq 2 && -n $1 && -n $2 ]] || die; cmd sudo mkfs.ext4 -F -L "$1" "$2"; }
fmt_fat32() { [[ $# -eq 2 && -n $1 && -n $2 ]] || die; cmd sudo mkfs.vfat -n"$1" "$2"; }

##
## 提示函数
##

# 给用户一个在继续或取消（退出此脚本）之间的选择
#  $1 标题
#  $2 文本
#
prompt_step()
{
  title="$1"
  msg="$2"
  unconditional="${3-}"
  if [[ ! ${unconditional-} && ${NOPROMPT:-} ]]; then
    echo -e "$msg"
    return 0
  fi
  zenity --title "$title" --question --ok-label "继续" --cancel-label "取消" --no-wrap --text "$msg"
  [[ $? = 0 ]] || exit 1
}

prompt_reboot()
{
  local msg=$1
  local mode="重启"
  [[ ${POWEROFF:-} ]] && mode="关机"

  prompt_step "操作成功" "${msg}\n\n选择继续立即$mode，或取消保留在修复镜像中。" "${REBOOTPROMPT:-}"
  [[ $? = 0 ]] || exit 1
  if [[ ${POWEROFF:-} ]]; then
    cmd systemctl poweroff
  else
    cmd systemctl reboot
  fi
}

##
## 修复函数
##

# 验证目标磁盘上的分区 - 至少确保类型和分区标签与预期的匹配。
#   $1 设备
#   $2 预期类型
#   $3 预期分区标签
#
verifypart()
{
  [[ $DOPARTVERIFY = 1 ]] || return 0
  TYPE="$(blkid -o value -s TYPE "$1" )"
  PARTLABEL="$(blkid -o value -s PARTLABEL "$1" )"
  if [[ ! $TYPE = "$2" ]]; then
    eerr "设备 $1 的类型是 $TYPE，但预期为 $2 - 无法继续。您可以尝试完全恢复。"
    sleep infinity ; exit 1
  fi

  if [[ ! $PARTLABEL = $3 ]] ; then 
    eerr "设备 $1 的标签是 $PARTLABEL，但预期为 $3 - 无法继续。您可以尝试完全恢复。"
    sleep infinity ; exit 2
  fi
}

# 替换设备 rootfs（btrfs 版本）。调用前必须冻结源。
#   $1 源设备
#   $2 目标设备
#
imageroot()
{
  local srcroot="$1"
  local newroot="$2"
  # 复制然后随机化目标 UUID - 这里要小心！复制 btrfs id 是有问题的
  cmd dd if="$srcroot" of="$newroot" bs=128M status=progress oflag=sync
  cmd btrfstune -f -u "$newroot"
  cmd btrfs check "$newroot"
}

# 在目标分区集中设置启动配置
#   $1 分区集名称
#
finalize_part()
{
  estat "完成安装分区 $1"
  cmd steamos-chroot --no-overlay --disk "$DISK" --partset "$1" -- mkdir /efi/SteamOS
  cmd steamos-chroot --no-overlay --disk "$DISK" --partset "$1" -- mkdir -p /esp/SteamOS/conf
  cmd steamos-chroot --no-overlay --disk "$DISK" --partset "$1" -- steamos-partsets /efi/SteamOS/partsets
  cmd steamos-chroot --no-overlay --disk "$DISK" --partset "$1" -- steamos-bootconf create --image "$1" --conf-dir /esp/SteamOS/conf --efi-dir /efi --set title "$1"
  cmd steamos-chroot --no-overlay --disk "$DISK" --partset "$1" -- grub-mkimage
  cmd steamos-chroot --no-overlay --disk "$DISK" --partset "$1" -- update-grub
}

##
## 主程序
##

onexit=()
exithandler() {
  for func in "${onexit[@]}"; do
    "$func"
  done
}
trap exithandler EXIT

# 检查目标磁盘是否存在
if [[ ! -e "$DISK" ]]; then
  eerr "$DISK 不存在 -- 未检测到 nvme 驱动器？"
  sleep infinity
  exit 1
fi

# 重新安装全新的 SteamOS 副本。
repair_steps()
{
  if [[ $writePartitionTable = 1 ]]; then
    estat "写入已知分区表"
    echo "$PARTITION_TABLE" | sfdisk "$DISK"

  elif [[ $writeOS = 1 || $writeHome = 1 ]]; then

    # 验证分区设置
    verifypart "$(diskpart $FS_ESP)" vfat esp
    verifypart "$(diskpart $FS_EFI_A)" vfat efi-A
    verifypart "$(diskpart $FS_EFI_B)" vfat efi-B
    verifypart "$(diskpart $FS_VAR_A)" ext4 var-A
    verifypart "$(diskpart $FS_VAR_B)" ext4 var-B
    verifypart "$(diskpart $FS_HOME)" ext4 home
  fi

  # 清除 var 分区（用户数据）
  if [[ $writeOS = 1 || $writeHome = 1 ]]; then
    estat "创建 var 分区"
    fmt_ext4  var  "$(diskpart $FS_VAR_A)"
    fmt_ext4  var  "$(diskpart $FS_VAR_B)"
  fi

  # 创建启动分区
  if [[ $writeOS = 1 ]]; then
    # 设置 ESP/EFI 启动分区
    estat "创建启动分区"
    fmt_fat32 esp  "$(diskpart $FS_ESP)"
    fmt_fat32 efi  "$(diskpart $FS_EFI_A)"
    fmt_fat32 efi  "$(diskpart $FS_EFI_B)"
  fi

  if [[ $writeHome = 1 ]]; then
    estat "创建 home 分区..."
    cmd sudo mkfs.ext4 -F -O casefold -T huge -L home "$(diskpart $FS_HOME)"
    estat "移除 home 分区上的预留块..."
    tune2fs -m 0 "$(diskpart $FS_HOME)"
  fi

  # 如果更新操作系统，为下次重启准备 BIOS 更新。
  if [[ $writeOS = 1 ]]; then
    estat "如有必要，为下次启动准备 BIOS 更新"
    biostool=/usr/bin/jupiter-biosupdate
    if [[ -n $VENDORED_BIOS_UPDATE && -d $VENDORED_BIOS_UPDATE ]]; then
      biostool="$VENDORED_BIOS_UPDATE"/jupiter-biosupdate
      export JUPITER_BIOS_DIR="$VENDORED_BIOS_UPDATE"
    fi

    fix_esp() {
      if [[ -n $mounted_esp ]]; then
        cmd umount -l /esp
        cmd umount -l /boot/efi
        mounted_esp=
      fi
    }
    onexit+=(fix_esp)
    einfo "在 /esp /boot/efi 上挂载新的 ESP/EFI 以进行 BIOS 准备"
    cmd mount "$(diskpart $FS_ESP)" /esp
    cmd mount "$(diskpart $FS_EFI_A)" /boot/efi
    mounted_esp=1

    if [[ ${FORCEBIOS:-} ]]; then
      "$biostool" --force || "$biostool"
    else
      "$biostool"
    fi

    fix_esp
  fi

  # 如果更新操作系统，执行控制器更新。OOBE 镜像不会在启动时自动更新控制器。
  if [[ $writeOS = 1 ]]; then
    estat "如有必要，更新控制器固件"
    controller_tool="/usr/bin/jupiter-controller-update"

    if [[ -n $VENDORED_CONTROLLER_UPDATE && -d $VENDORED_CONTROLLER_UPDATE ]]; then
      controller_tool="$VENDORED_CONTROLLER_UPDATE"/jupiter-controller-update
      export JUPITER_CONTROLLER_UPDATE_FIRMWARE_DIR="$VENDORED_CONTROLLER_UPDATE"
    fi

    JUPITER_CONTROLLER_UPDATE_IN_OOBE=1 "$controller_tool"
  fi

  if [[ $writeOS = 1 ]]; then
    rootdevice="$(findmnt -n -o source / )"
    if [[ -z $rootdevice || ! -e $rootdevice ]]; then
      eerr "找不到 USB 安装程序根目录 -- USB 集线器问题？"
      sleep infinity
      exit 1
    fi

    estat "冻结 rootfs"
    unfreeze() { fsfreeze -u /; }
    onexit+=(unfreeze)
    cmd fsfreeze -f /

    estat "成像操作系统分区 A"
    imageroot "$rootdevice" "$(diskpart $FS_ROOT_A)"
  
    estat "成像操作系统分区 B"
    imageroot "$rootdevice" "$(diskpart $FS_ROOT_B)"
  
    estat "完成启动配置"
    finalize_part A
    finalize_part B
    estat "完成 EFI 系统分区"
    cmd steamos-chroot --no-overlay --disk "$DISK" --partset A -- steamcl-install --flags restricted --force-extra-removable
  fi
}

# 进入设备上的主操作系统分区集
chroot_primary()
{
  partset=$( steamos-chroot --no-overlay --disk "$DISK" --partset "A" -- steamos-bootconf selected-image )

  estat "进入 $partset 分区集的 chroot。"
  estat "您可以在这里进行任何需要的更改，完成后退出。"

  # FIXME etc 覆盖目录在全新安装时可能不存在

  cmd steamos-chroot --disk "$DISK" --partset "$partset"
}

# 返回清理状态（并回显当前完成百分比）
# 0 : 准备清理
# 1 : 清理进行中（回显当前百分比）
# 2 : 驱动器不支持清理
#
get_sanitize_progress()
{
  status=$(nvme sanitize-log "${DISK}" | grep "(SSTAT)" | grep -oEi "(0x)?[[:xdigit:]]+$") || return 2
  [[ $(( status % 8 )) -eq 2 ]] || return 0

  progress=$(nvme sanitize-log "${DISK}" | grep "(SPROG)" | grep -oEi "(0x)?[[:xdigit:]]+$") || return 2
  echo "清理进度: $(( ( progress * 100 )/ 65535 ))%"
  return 1
}

# 调用 nvme 清理
sanitize_all()
{
  sres=0
  get_sanitize_progress || sres=$?
  case $sres in
    0)
      echo
      echo "警告！"
      echo
      echo "此操作将不可逆地清除 ${DISK} 上的 *所有* 用户数据"
      echo "暂停五秒钟，以防您不是有意执行此操作..."
      sleep 5
      echo "好的，开始执行。正在清理 ${DISK}："
      nvme sanitize -a 2 "${DISK}"
      echo "清理操作已启动。"
      ;;
    1) echo "NVME 清理操作已在进行中。"
      ;;
    2) # 使用 NVME 安全格式
      nvme format "${DISK}" -n 1 -s 1 -r
      return 0
      ;;
    *) echo "来自 sanitize-log 的意外结果"
      return $sres
      ;;
  esac

  while ! get_sanitize_progress ; do
    sleep 5
  done

  echo "... 清理完成。"
}

# 打印目标的列表
help()
{
  readvar HELPMSG << EOD
此工具可用于重新安装或修复您的 SteamOS 安装

可能的目标：
    all : 永久销毁设备上的所有数据，并（重新）安装 SteamOS。
    system : 修复/重新安装设备系统分区上的 SteamOS，保留用户数据分区。
    home : 重新格式化设备的 /home 和 /var 分区，从设备中删除游戏和用户数据。
    chroot : chroot 进入主 SteamOS 分区集。
    sanitize : 执行 NVME 清理操作。
EOD
  emsg "$HELPMSG"
  if [[ "$EUID" -ne 0 ]]; then
    eerr "请以 root 身份运行。"
    exit 1
  fi
}

[[ "$EUID" -ne 0 ]] && help

writePartitionTable=0
writeOS=0
writeHome=0

case "${1-help}" in
all)
  prompt_step "擦除设备并安装 SteamOS" "此操作将擦除并在此设备上（重新）安装 SteamOS。\n这将永久销毁您设备上的所有数据。\n\n这无法撤销。\n\n只有在您希望擦除并重新安装此设备时才选择继续。"
  writePartitionTable=1
  writeOS=1
  writeHome=1
  sanitize_all
  repair_steps
  prompt_reboot "重新成像完成。"
  ;;
system)
  prompt_step "修复 SteamOS" "此操作将修复设备上的 SteamOS 安装，同时尝试保留您的游戏和个人内容。\n系统自定义可能会丢失。\n\n选择继续在您的设备上重新安装 SteamOS。"
  writeOS=1
  repair_steps
  prompt_reboot "SteamOS 重新安装完成。"
  ;;
home)
  prompt_step "删除本地用户数据" "此操作将重新格式化您设备上的 home 分区。\n这将销毁下载的游戏和所有个人内容，包括系统配置。\n\n此操作无法撤销。\n\n选择继续重新格式化所有用户 home 分区。"
  writeHome=1
  repair_steps
  prompt_reboot "用户分区已重新格式化。"
  ;;
chroot)
  chroot_primary
  ;;
sanitize)
  prompt_step "清除并清理 NVME 磁盘" "此操作将对主驱动器启动 NVME 清理，不可逆地删除所有用户数据。\n\n此操作无法撤销。\n\n只有在您要从当前设备的主驱动器中删除所有数据时才选择继续。"
  sanitize_all
  ;;
*)
  help
  ;;
esac 
